#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rclpy
from rclpy.node import Node
from rclpy.action import ActionServer, GoalResponse, CancelResponse
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup

from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
from pyzbar import pyzbar
import time
import yaml
import os
from ament_index_python.packages import get_package_share_directory

from qr_recognition.action import RecognizeQR


class QRRecognitionServer(Node):
    """二维码识别Action服务器节点"""
    
    def __init__(self):
        super().__init__('qr_recognition_server')
        
        self.get_logger().info('初始化二维码识别服务器...')
        
        # CV Bridge用于转换ROS图像消息
        self.bridge = CvBridge()
        
        # 存储最新的图像
        self.current_image = None
        self.image_lock = False
        
        # 加载配置文件
        self.config = self.load_config()
        self.get_logger().info('已加载默认数据配置文件')
        
        # 创建回调组
        self.callback_group = ReentrantCallbackGroup()
        
        # 订阅图像话题
        self.image_sub = self.create_subscription(
            Image,
            '/left_cam/image_raw',
            self.image_callback,
            10,
            callback_group=self.callback_group
        )
        
        # 创建Action服务器
        self._action_server = ActionServer(
            self,
            RecognizeQR,
            'recognize_qr',
            execute_callback=self.execute_callback,
            goal_callback=self.goal_callback,
            cancel_callback=self.cancel_callback,
            callback_group=self.callback_group
        )
        
        self.get_logger().info('二维码识别服务器已启动，等待Action请求...')
    
    def load_config(self):
        """加载配置文件"""
        try:
            # 尝试从package的config目录加载
            package_share_dir = get_package_share_directory('qr_recognition')
            config_path = os.path.join(package_share_dir, 'config', 'default_data.yaml')
            
            # 如果在install目录找不到，尝试从源码目录加载
            if not os.path.exists(config_path):
                # 获取当前脚本所在目录的上一级目录
                script_dir = os.path.dirname(os.path.abspath(__file__))
                package_dir = os.path.dirname(script_dir)
                config_path = os.path.join(package_dir, 'config', 'default_data.yaml')
            
            self.get_logger().info(f'配置文件路径: {config_path}')
            
            with open(config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            
            return config
        except Exception as e:
            self.get_logger().error(f'加载配置文件失败: {str(e)}，使用内置默认值')
            # 返回内置的默认配置
            return {
                'type1': {
                    'fruits': ['苹果', '梨子', '苹果', '梨子', '苹果', '梨子', '苹果', '梨子']
                },
                'type2': {
                    'vegetables': ['洋葱', '南瓜', '西红柿', '辣椒', '洋葱', '南瓜', '西红柿', '辣椒'],
                    'numbers': [1, 2, 3, 4, 8, 9, 12, 11]
                },
                'name_mapping': {
                    '苹果': 1, '梨子': 2,
                    '洋葱': 1, '南瓜': 2, '西红柿': 3, '辣椒': 4
                }
            }
    
    def image_callback(self, msg):
        """图像话题回调函数"""
        try:
            # 将ROS图像消息转换为OpenCV图像
            self.current_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            self.get_logger().error(f'图像转换失败: {str(e)}')
    
    def goal_callback(self, goal_request):
        """处理新的Action请求"""
        self.get_logger().info(f'收到识别请求，类型参数: {goal_request.type}')
        return GoalResponse.ACCEPT
    
    def cancel_callback(self, goal_handle):
        """处理取消请求"""
        self.get_logger().info('收到取消请求')
        return CancelResponse.ACCEPT
    
    def execute_callback(self, goal_handle):
        """执行二维码识别任务"""
        self.get_logger().info('开始执行二维码识别...')
        
        # 获取请求参数
        request_type = goal_handle.request.type
        
        # 发送反馈
        feedback_msg = RecognizeQR.Feedback()
        feedback_msg.status = '等待图像数据...'
        goal_handle.publish_feedback(feedback_msg)
        
        # 等待图像数据
        retry_count = 0
        max_retries = 50  # 最多等待5秒
        while self.current_image is None and retry_count < max_retries:
            time.sleep(0.1)
            retry_count += 1
        
        if self.current_image is None:
            self.get_logger().warn('未能获取图像数据，将使用配置的默认数据')
            # 直接使用默认数据，不再abort
            fruit_array, number_array = self.get_default_data(request_type)
            
            feedback_msg.status = '无图像数据，使用默认值'
            goal_handle.publish_feedback(feedback_msg)
            
            # 设置成功并返回结果
            goal_handle.succeed()
            result = RecognizeQR.Result()
            result.fruit_array = fruit_array
            result.number_array = number_array
            
            self.get_logger().info(f'返回默认蔬菜/水果编码: {fruit_array}')
            if number_array:
                self.get_logger().info(f'返回默认数字数据: {number_array}')
            
            return result
        
        feedback_msg.status = '正在识别二维码...'
        goal_handle.publish_feedback(feedback_msg)
        
        # 持续1秒进行识别，只要有一次成功就返回
        qr_text = None
        start_time = time.time()
        recognition_duration = 1.0  # 识别持续时间（秒）
        attempt_count = 0
        
        while (time.time() - start_time) < recognition_duration:
            # 获取当前最新的图像
            if self.current_image is not None:
                # 尝试识别二维码
                qr_text = self.detect_qr_code(self.current_image)
                attempt_count += 1
                
                # 如果识别成功，立即跳出循环
                if qr_text is not None:
                    self.get_logger().info(f'第{attempt_count}次尝试成功识别二维码！')
                    break
            
            # 短暂等待后继续尝试
            time.sleep(0.05)
        
        # 如果在1秒内都没有识别成功，使用配置的默认数据
        if qr_text is None:
            self.get_logger().warn(f'尝试{attempt_count}次后未检测到二维码，使用配置的默认数据')
            # 直接获取默认数据
            fruit_array, number_array = self.get_default_data(request_type)
        else:
            self.get_logger().info(f'识别到的二维码内容: {qr_text}')
            # 解析识别到的二维码内容
            fruit_array, number_array = self.parse_qr_text(qr_text, request_type)
        
        feedback_msg.status = '识别完成'
        goal_handle.publish_feedback(feedback_msg)
        
        self.get_logger().info(f'蔬菜/水果编码: {fruit_array}')
        if number_array:
            self.get_logger().info(f'数字数据: {number_array}')
        
        # 设置成功并返回结果
        goal_handle.succeed()
        
        result = RecognizeQR.Result()
        result.fruit_array = fruit_array
        result.number_array = number_array
        
        return result
    
    def detect_qr_code(self, image):
        """
        使用pyzbar检测并解码二维码
        
        Args:
            image: OpenCV图像
            
        Returns:
            str: 二维码内容，如果未检测到则返回None
        """
        try:
            # 使用pyzbar解码二维码
            decoded_objects = pyzbar.decode(image)
            
            # 如果检测到二维码，返回第一个二维码的数据
            if decoded_objects:
                # 获取第一个二维码的数据并解码为字符串
                qr_data = decoded_objects[0].data.decode('utf-8')
                return qr_data
            else:
                return None
                
        except Exception as e:
            self.get_logger().error(f'二维码检测失败: {str(e)}')
            return None
    
    def get_default_data(self, request_type):
        """
        从配置文件获取默认数据
        
        Args:
            request_type: 请求类型（1或2）
            
        Returns:
            tuple: (fruit_array, number_array)
        """
        fruit_array = []
        number_array = []
        
        try:
            name_mapping = self.config.get('name_mapping', {})
            
            if request_type == 1:
                # type1: 水果
                fruits = self.config.get('type1', {}).get('fruits', [])
                for fruit_name in fruits:
                    if fruit_name in name_mapping:
                        fruit_array.append(name_mapping[fruit_name])
                        
            elif request_type == 2:
                # type2: 蔬菜和数字
                vegetables = self.config.get('type2', {}).get('vegetables', [])
                for veg_name in vegetables:
                    if veg_name in name_mapping:
                        fruit_array.append(name_mapping[veg_name])
                
                # 获取数字数组
                number_array = self.config.get('type2', {}).get('numbers', [])
        
        except Exception as e:
            self.get_logger().error(f'获取默认数据失败: {str(e)}')
        
        return fruit_array, number_array
    
    def parse_qr_text(self, text, request_type):
        """
        解析二维码文本，转换为数组
        支持多种格式：
        1. type1 - 水果名称（每行一个）：苹果 -> 1, 梨子 -> 2
        2. type2 - 蔬菜名称+数字（前8行蔬菜名称，第9行数字）
           - 蔬菜：洋葱 -> 1, 南瓜 -> 2, 西红柿 -> 3, 辣椒 -> 4
           - 数字：1,2,3,4,8,9,12,11
        
        Args:
            text: 二维码文本内容
            request_type: 请求类型（1或2）
            
        Returns:
            tuple: (fruit_array, number_array)
        """
        fruit_array = []
        number_array = []
        
        if not text:
            return fruit_array, number_array
        
        text = text.strip()
        lines = text.split('\n')
        
        # 从配置文件获取名称到编号的映射
        name_mapping = self.config.get('name_mapping', {})
        
        if request_type == 1:
            # type1: 只解析水果名称
            for line in lines:
                item = line.strip()
                if item in name_mapping:
                    fruit_array.append(name_mapping[item])
            return fruit_array, number_array
        
        elif request_type == 2:
            # type2: 前8行是蔬菜名称，第9行是数字
            for i, line in enumerate(lines):
                line = line.strip()
                
                # 前8行：解析蔬菜名称
                if i < 8:
                    if line in name_mapping:
                        fruit_array.append(name_mapping[line])
                # 第9行及之后：解析数字（逗号分隔）
                else:
                    if ',' in line:
                        try:
                            numbers = line.split(',')
                            for num in numbers:
                                num = num.strip()
                                if num.isdigit():
                                    number_array.append(int(num))
                        except Exception as e:
                            self.get_logger().error(f'解析数字格式失败: {str(e)}')
            
            return fruit_array, number_array
        
        else:
            # 其他类型：尝试通用解析
            for line in lines:
                item = line.strip()
                if item in name_mapping:
                    fruit_array.append(name_mapping[item])
            return fruit_array, number_array


def main(args=None):
    rclpy.init(args=args)
    
    qr_recognition_server = QRRecognitionServer()
    
    # 使用多线程执行器
    executor = MultiThreadedExecutor()
    executor.add_node(qr_recognition_server)
    
    try:
        executor.spin()
    except KeyboardInterrupt:
        pass
    finally:
        qr_recognition_server.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()

